"""
缓存模型响应消息
"""
from langchain_community.cache import SQLiteCache, SQLAlchemyCache, FullLLMCache
from langchain_core.caches import InMemoryCache
from langchain_core.globals import set_llm_cache
from sqlalchemy import create_engine, MetaData, func
from sqlalchemy.orm import declarative_base
from sqlalchemy.sql.schema import Table, Column
from sqlalchemy.sql.sqltypes import String, Text, DateTime, Integer

from models import get_ds_model_client

client = get_ds_model_client()

# 基于内存缓存
set_llm_cache(InMemoryCache())
print(client.invoke("Tell me a joke, 50 character"))

# 基于文件缓存
set_llm_cache(SQLiteCache(database_path="../data/cache/langchain.db"))
print(client.invoke("hello world"))

# 基于Mysql缓存
engine = create_engine('mysql+pymysql://root:root@192.168.46.128:3306/test')
Base = declarative_base()
class CustomLLMCache(Base):
    __tablename__ = 'llm_cache'
    idx = Column(Integer, primary_key=True, autoincrement=True)
    prompt = Column(Text)
    llm = Column(Text)
    response = Column(Text)
    created_at = Column(DateTime, default=func.now())
set_llm_cache(SQLAlchemyCache(engine, cache_schema=CustomLLMCache))
print(client.invoke("hello world"))

